# This file is part of Neuroinfo Toolkit.
#
# Neuroinfo Toolkit is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Neuroinfo Toolkit is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Neuroinfo Toolkit.  If not, see <http://www.gnu.org/licenses/>.

import gzip
import struct
import neuro.filesystem as filesystem
import neuro.system as system
import neuro.strings as strings
from neuro.models.resource import Resource
from neuro.command import Command
from neuro.command import CommandFailedException
from neuro.exceptions import IllegalArgumentException
from neuro.exceptions import FileFormatException
from neuro.exceptions import IOException

class Nifti(Resource):
	'''
	NIFTI image
	
	.. note:: For simplicity sake, many of these class methods outsource to FSL utilities. In the future, these implementations may be rewritten in favor of a more native solution.
	'''
	
	def __init__(self, input):
		'''
		Constructor ::
		
			>>> nifti = new Nifti()
			>>> nifti = new Nifti("/path/to/nifti.nii.gz")
		
		:param input: Input file name or :class:`~neuro.models.resource.Resource`
		:type input: str, :class:`~neuro.models.resource.Resource`
		'''
		Resource.__init__(self, input)

		self._filename = None
		self._tr = None;
		self._x = None;
		self._y = None;
		self._z = None;
		self._t = None;
		self._numDim = None;
		self._orientation = None;
		self._minIntensity = None;
		self._maxIntensity = None;
		
		## --- import mixed input variable
		if(isinstance(input, basestring)):
			self._importFile(input)
		elif(isinstance(input, Resource)):
			self._importResource(input)
		else:
			raise IllegalArgumentException("Input must be an instance of str or Resource")
	
	def _importFile(self, filename):
		'''
		Import NIFTI file
			
		:param filename: File name
		:type filename: str
		:raises: :class:`~neuro.exceptions.IOException`,
				 :class:`~neuro.exceptions.FileFormatException`
		'''
		filename = filesystem.canonical(filename)

		## --- check if input file is readable
		if(not filesystem.isReadable(filename)):
			raise IOException(IOException.READ, filename)

		if(not Nifti.checkFormat(filename)):
			raise FileFormatException(filename)

		self._filename = filename
	
	@staticmethod   
	def checkFormat(filename):
		'''
		Check file format validity ::
		
			Nifti.checkFormat("/path/to/nifti.nii.gz")
			
		:param filename: File name
		:type filename: str
		:returns: File format validity
		:rtype: bool
		:raises: :class:`~neuro.exceptions.IOException`,
				 :class:`~neuro.exceptions.FileFormatException`
		'''   
		Resource.checkFormat(filename)

		filename = filesystem.canonical(filename)

		ext = filesystem.getExtension(filename)
		
		if(ext == ".gz"):
			h = gzip.open(filename, 'rb')
		else:
			h = open(filename, 'rb')
			
		try:
			dat = struct.unpack("i", h.read(4))
			
			if(dat[0] == 348):
				return True
			else:
				return False
		except FileFormatException, ffe:
			raise
		except Exception, e:
			return False
		
	def getX(self):
		'''
		Retrieve X dimension ::
		
			>>> nifti.getX()
			72
		
		:returns: Volume X dimension (in voxels)
		:rtype: int
		:raises: :class:`~neuro.command.CommandFailedException`,
				 :class:`~neuro.exceptions.CommandNotFoundException`
		''' 
		if(self._x == None):
			system.check("fslhd", "FMRIB Software Library")
			
			command = Command('fslhd -x ' + self._filename + ' | grep nx | sed -re "s@.*nx\s=\s\'(.*)\'@\\1@\"')
			command.execute()
			
			x = command.getStdout()
				
			if(x == ""):
				raise CommandFailedException(command)
			
			self._x = x

		return self._x
	
	def getY(self):
		'''
		Retrieve Y dimension ::
		
			>>> nifti.getY()
			72
		
		:returns: Volume Y dimension (in voxels)
		:rtype: int
		:raises: :class:`~neuro.command.CommandFailedException`,
				 :class:`~neuro.exceptions.CommandNotFoundException`
		'''
		if(self._y == None):
			system.check("fslhd", "FMRIB Software Library")

			command = Command('fslhd -x ' + self._filename + ' | grep ny | sed -re "s@.*ny\s=\s\'(.*)\'@\\1@\"')
			command.execute()

			y = command.getStdout()

			if(y == ""):
				raise CommandFailedException(command)

			self._y = y
				
		return self._y
	
	def getZ(self):
		'''
		Retrieve Z dimension ::
		
			>>> nifti.getZ()
			47
		
		:returns: Volume Z dimension (in voxels)
		:rtype: int
		:raises: :class:`~neuro.command.CommandFailedException`,
				 :class:`~neuro.exceptions.CommandNotFoundException`
		'''
		if(self._z == None):
			system.check("fslhd", "FMRIB Software Library")

			command = Command('fslhd -x ' + self._filename + ' | grep nz | sed -re "s@.*nz\s=\s\'(.*)\'@\\1@\"')
			command.execute()

			z = command.getStdout()

			if(z == ""):
				raise CommandFailedException(command)

			self._z = z

		return self._z
	
	def getT(self):
		'''
		Retrieve number of timepoints ::
		
			>>> nifti.getT()
			124
			
		:rtype: int
		:raises: :class:`~neuro.command.CommandFailedException`,
				 :class:`~neuro.exceptions.CommandNotFoundException`
		'''
		if(self._t == None):
			system.check("fslhd", "FMRIB Software Library")

			dim = self.getNumDimensions()

			if(dim == 3):
				self._t = 1
			elif(dim == 4):
				command = Command('fslhd -x ' + self._filename + ' | grep "nt =" | sed -re "s@.*nt\s=\s\'(.*)\'@\\1@\"')
				command.execute()

				t = command.getStdout()

				if(t == ""):
					raise CommandFailedException(command)

				self._t = t
			else:
				raise InternalException("Unexpected number of dimensions: " + dim)

		return self._t
	
	def getNumVolumes(self):
		'''
		Retrieve number of volumes/timepoints. Same as :class:`Nifti.getT()`::
		
			>>> nifti.getNumVolumes()
			124
			
		:returns: Number of volumes/timepoints
		:rtype: int
		:raises: :class:`~neuro.command.CommandFailedException`,
				 :class:`~neuro.exceptions.CommandNotFoundException`
		'''
		return self.getT()

	def getNumDimensions(self):
		'''
		Retrieve number of dimensions e.g. 3-D, 4-D ::
		
			>>> dims = nifti.getNumDimensions()
			4
		
		:returns: Number of dimensions
		:rtype: int
		:raises: :class:`~neuro.command.CommandFailedException`,
				 :class:`~neuro.exceptions.CommandNotFoundException`
		'''
		## --- check for URI 
		if(self._uri == None):
			raise MissingFieldException("URI field is not for this object")
		
		if(self._numDim == None):
			system.check("fslhd", "FMRIB Software Library")

			command = Command('fslhd -x ' + self._filename + ' | grep ndim | sed -re "s@.*ndim\s=\s\'(.*)\'@\\1@\"')
			command.execute()

			numDim = command.getStdout()

			if(numDim == ""):
				raise CommandFailedException(command)

			self._numDim = int(numDim)
		
		return self._numDim

	def getRepetitionTime(self):
		'''
		Retrieve repetition time (TR). Not defined in NIFTI format ::
		
			>>> nifti.getRepetitionTime()
			3000.0
		
		:returns: Repetition time (TR)
		:rtype: float
		'''
		return None
	
	def getOrientation(self):
		'''
		Retrieve orientation ::
		
			>>> nifti.getOrientation()
			'radiological'
		
		:returns: Image left/right orientation
		:rtype: str
		:raises: :class:`~neuro.command.CommandFailedException`,
				 :class:`~neuro.exceptions.CommandNotFoundException`
		'''
		if(self._orientation == None):
			system.check("fslorient", "FMRIB Software Library")

			command = Command('fslorient -getorient ' + self._filename)
			command.execute()

			orientation = command.getStdout()

			if(orientation == ""):
				raise CommandFailedException(command)

			self._orientation = orientation.lower()
		
		return self._orientation
	
	def setOrientation(self, orientation):
		'''		
		.. warning:: Realize that this method will swap both the header *and* voxel data.
		
		Sets the orientation of a NIFTI file ::
		
			>>> nifti.setOrientation("neurological")
			'neurological'
			
		:param orientation: "radiological" or "neurological"
		:type orientation: str
		:returns: New orientation
		:rtype: str
		:raises: :class:`~neuro.command.CommandFailedException`,
				 :class:`~neuro.exceptions.CommandNotFoundException`
		'''
		if(not isinstance(orientation, basestring)):
			raise IllegalArgumentException("Orientation must be a string")
		
		orgOrientation = self.getOrientation()
		orientation = orientation.lower()
		
		if(orientation == "radiological"):
			if(orgOrientation != "radiological"):
				self.invertAxis("x")
		elif(orientation == "neurological"):
			if(orgOrientation != "neurological"):
				self.invertAxis("x")
		else:
			raise IllegalArgumentException("Orientation must be either \"radiological\" or \"neurological\"")
		
		return self._orientation
			
	def flipOrientation(self):
		'''
		.. warning:: Realize that this method will swap both the header *and* voxel data.
		
		Flip orientation ::
		
			>>> nifti.flipOrientation()
			'radiological'
		
		:returns: New orientation
		:rtype: str
		:raises: :class:`~neuro.command.CommandFailedException`,
				 :class:`~neuro.exceptions.CommandNotFoundException`
		'''
		self.invertAxis("x")
		
		return self._orientation

	def getMinIntensity(self):
		'''
		Retrieve minimum intensity ::
		
			>>> nifti.getMinIntensity()
			0.000000
			
		:returns: Minimim voxel intensity
		:rtype: float
		:raises: :class:`~neuro.command.CommandFailedException`,
				 :class:`~neuro.exceptions.CommandNotFoundException`
		'''	
		if(self._minIntensity == None):
			system.check("fslstats", "FMRIB Software Library")

			command = Command('fslstats ' + self._filename + ' -R | cut -d\' \' -f1')
			command.execute()

			intensity = command.getStdout()

			if(intensity == ""):
				raise CommandFailedException(command)

			## --- clear and regenerate orientation
			self._minIntensity = float(intensity)
		
		return self._minIntensity
	
	def getMaxIntensity(self):
		'''
		Retrieve maximum intensity ::
		
			>>> nifti.getMaxIntensity()
			3911.000000
			
		:returns: Maximum voxel intensity
		:rtype: float
		:raises: :class:`~neuro.command.CommandFailedException`,
				 :class:`~neuro.exceptions.CommandNotFoundException`
		'''
		if(self._maxIntensity == None):
			system.check("fslstats", "FMRIB Software Library")

			command = Command('fslstats ' + self._filename + ' -R | cut -d\' \' -f2')
			command.execute()

			intensity = command.getStdout()

			if(intensity == ""):
				raise CommandFailedException(command)

			## --- clear and regenerate orientation
			self._maxIntensity = float(intensity)
		
		return self._maxIntensity
	
	def invertAxis(self, dim):
		'''
		.. warning:: Realize that this method will swap both the header *and* voxel data.
		
		Invert a dimension ::
		
			>>> nifti.invertAxis("x")
			
		:param dim: Axis to invert "x", "y", or "z"
		:type dim: str
		:raises: :class:`~neuro.command.CommandFailedException`,
				 :class:`~neuro.exceptions.CommandNotFoundException`
		'''	
		if(not isinstance(dim, basestring)):
			raise IllegalArgumentException("Dimension must be an instance of str")
		
		dim = dim.lower()
			
		system.check("fslswapdim", "FMRIB Software Library")
		
		if(dim == "x"):
			system.check("fslorient", "FMRIB Software Library")
			
			## --- flip the actual pixels
			flipPayload = Command("fslswapdim " + self._filename + " -x y z " + self._filename)
			flipPayload.execute()

			## --- flip the header
			flipHeader = Command("fslorient -swaporient " + self._filename)
			flipHeader.execute()
			
			## --- reset orientation field
			self._orientation = None
			self.getOrientation()
		elif(dim == "y"):
			command = Command("fslswapdim " + self._filename + " x -y z " + self._filename)
			command.execute()
		elif(dim == "z"):
			command = Command("fslswapdim " + self._filename + " x y -z " + self._filename)
			command.execute()
		else:
			raise IllegalArgumentException("Dimension must be \"x\", \"y\", or \"z\"")
		
	def getTimeseries(self, x, y, z):
		'''
		Retrieve voxel timeseries ::
		
			>>> nifti.getTimeSeries(1,2,3)
			
			
		:param x:
		:type x: int
		:param y:
		:type y: int
		:param z:
		:type z: int
		:returns: Voxel timeseries
		:rtype: list
		:raises: :class:`NiftiDimensionException`,
				 :class:`~neuro.command.CommandFailedException`
		'''
		if(x == None):
			raise IllegalArgumentException("X cannot be null")

		if(y == None):
			raise IllegalArgumentException("Y cannot be null")
		
		if(z == None):
			raise IllegalArgumentException("Z cannot be null")
		
		if(not isinstance(x, int)):
			raise IllegalArgumentException("X must be an instance of int")
		
		if(x < 0 or x >= self.getX()):
			raise NiftiDimensionException(NiftiDimensionException.X_RANGE, self, x)
		
		if(y < 0 or y >= self.getY()):
			raise NiftiDimensionException(NiftiDimensionException.Y_RANGE, self, y)
		
		if(y < 0 or y >= self.getZ()):
			raise NiftiDimensionException(NiftiDimensionException.Z_RANGE, self, z)
		
		system.check("fslmeants", "FMRIB Software Library")

		command = Command("fslmeants -i " + self._filename + " -c " + str(x) + " " + str(y) + " " + str(z))
		command.execute()

		if(command.getStdout() == ""):
			raise CommandFailedException(command)

		output = command.getStdout()
		output = output.split("\n")
			
		if(len(output) == 0):
			raise InternalException("Splitting timeseries ended up with an empty result")
			
		return map(float, output)
	
	def getExtension(self):
		'''
		Retrieve NIFTI file extension
		
			>>> nifti = new Nifti("/path/to/nifti.nii.gz")
			>>> nifti.getExtension()
			'.nii.gz'
		
		:returns: NIFTI file extension
		:rtype: str
		'''
		extension = strings.regex("^.*(\.nii$|\.nii\.gz$)", self._filename)
		
		if(extension):
			return extension[0]
		else:
			return ""

class NiftiDimensionException(BaseException):
	'''
	NIFTI-1 data dimension exception
	'''
	X_RANGE=1
	Y_RANGE=2
	Z_RANGE=3
	T_RANGE=4

	def __init__(self, type, nifti, dim):
		'''
		Constructor

		:param type: Type of exception
		:type type: str
		:param nifti: :class:`Nifti` object
		:type nifti: :class:`Nifti`
		:param dim: Dimension
		:type dim: str
		'''
		if(not isinstance(type, int)):
			raise IllegalArgumentException("Exception type must be an instance of int")
		elif(not isinstance(nifti, Nifti)):
			raise IllegalArgumentException("NIFTI input parameter must be an instance of Nifti")
		elif(not isinstance(dim, int)):
			raise IllegalArgumentException("NIFTI dimension must be an instance of int")

		if(type < 1 or type > 4):
			raise IllegalArgumentException("Exception type must be NiftiDimensionException.X_RANGE, NiftiDimensionException.Y_RANGE, NiftiDimensionException.Z_RANGE, or NiftiDimensionException.T_RANGE")

		self._type = type
		self._nifti = nifti
		self._dim = dim

	def getType(self):
		'''
		Get exception type

		:rtype: int
		'''
		return self._type

	def getNifti(self):
		'''
		Get Nifti that triggered this exception

		:rtype: :class:`Nifti`
		'''
		return self._nifti

	def getDimension(self):
		'''
		Get requested dimension that triggered this exception

		:rtype: int
		'''
		return self._dim

	def getMessage(self):
		'''
		Get custom message

		:rtype: str
		'''
		if(self._type == NiftiDimensionException.X_RANGE):
			return "Requested point " + str(self._dim) + " along X dimension of " + str(self._nifti.getX()) + " possible points (voxels)"
		elif(self._type == NiftiDimensionException.Y_RANGE):
			return "Requested point " + str(self._dim) + " along Y dimension of " + str(self._nifti.getY()) + " possible points (voxels)"
		elif(self._type == NiftiDimensionException.Z_RANGE):
			return "Requested point " + str(self._dim) + " along Z dimension of " + str(self._nifti.getZ()) + " possible points (voxels)"
		elif(self._type == NiftiDimensionException.T_RANGE):
			return "Requested point " + str(self._dim) + " along T dimension of " + str(self._nifti.getT()) + " possible points (time)"